from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

import torch.nn as nn

from continuallearning.interfaces import (
    TaskAwareHookManagerInterface,
    TaskAwareHookInterface,
)

# 替换标准logging库为自定义日志系统
from continuallearning.utils.logging import get_logger
from continuallearning.registry import HOOK_MANAGER_REGISTRY

# 使用自定义日志系统创建记录器
logger = get_logger(__name__)


@HOOK_MANAGER_REGISTRY.register("base")
class AbstractHookManager(TaskAwareHookManagerInterface, ABC):
    """
    Abstract base class for hook managers.

    Hook managers are responsible for coordinating the lifecycle of multiple hooks
    across a model, providing centralized management and access to hook parameters.

    This class defines the core interface but delegates implementation details
    to subclasses.
    """

    _current_task: int = -1

    def __init__(self, *args, **kwargs):
        """Initialize the hook manager."""
        self._initialized = False
        self._is_attached = False
        # hooks that are attached to different modules
        self._hooks: Dict[str, TaskAwareHookInterface] = {}
        # Centralized task management
        self._registered_tasks: List[int] = []
        self._task_configs: Dict[int, Optional[Dict[str, Any]]] = {}

    @property
    def is_initialized(self) -> bool:
        """Check if the hook manager is initialized."""
        return self._initialized

    @property
    def is_attached(self) -> bool:
        """Check if hooks are currently attached."""
        return self._is_attached

    @property
    def hooks(self) -> Dict[str, TaskAwareHookInterface]:
        """
        Get the dictionary of hooks managed by this hook manager.

        Returns:
            Dict[str, HookInterface]: Dictionary mapping hook identifiers to hook objects
        """
        return self._hooks

    def _initialize_resources(self) -> None:
        """
        Initialize required resources for hook operations.

        This method is called before adding hooks to ensure all necessary
        resources are properly configured. Subclasses can override this method
        to allocate additional resources needed for specific hook types.
        """
        if self._initialized:
            return

        self._initialized = True

    @abstractmethod
    def add_hooks(self) -> None:
        """Add hooks to the model. set is_attached status"""
        self._is_attached = True

    @abstractmethod
    def remove_hooks(self) -> None:
        """Remove all hooks from the model."""
        pass

    def attach(self) -> "TaskAwareHookManagerInterface":
        """
        Attach all hooks to the model.

        Returns:
            HookManager: Self, for method chaining
        """
        if not self._initialized:
            self._initialize_resources()

        if not self._is_attached:
            self.add_hooks()

        return self

    def detach(self) -> None:
        """Detach all hooks from the model."""
        if self._is_attached:
            self.remove_hooks()
            self._is_attached = False

    def __enter__(self) -> "TaskAwareHookManagerInterface":
        """
        Enter context manager, attaching hooks.

        Returns:
            HookManager: Self for use in with statements
        """
        self.attach()
        return self

    def __exit__(self, exc_type=None, exc_val=None, exc_tb=None) -> None:
        """
        Exit context manager, detaching hooks.

        Args:
            exc_type: Exception type if an exception was raised
            exc_val: Exception value if an exception was raised
            exc_tb: Exception traceback if an exception was raised
        """
        self.detach()

    # hook interaction
    def _apply_to_hooks(self, func_name: str, *args, **kwargs) -> None:
        """
        Apply a method to all task-aware hooks.

        Args:
            func_name: Name of the method to call
            *args, **kwargs: Arguments to pass to the method
        """
        for hook in self._hooks.values():
            method = getattr(hook, func_name)
            method(*args, **kwargs)

    def _collect_from_hooks(self, func_name: str, *args, **kwargs) -> List[Any]:
        """
        Collect results from calling a method on all hooks.

        Args:
            func_name: Name of the method to call
            *args, **kwargs: Arguments to pass to the method

        Returns:
            List[Any]: Flattened list of results from all hooks
        """
        results = []
        for hook in self._hooks.values():
            method = getattr(hook, func_name)
            hook_results = method(*args, **kwargs)
            if isinstance(hook_results, list):
                results.extend(hook_results)
            else:
                results.append(hook_results)
        return results

    def get_parameters(self) -> List[nn.Parameter]:
        """
        Get all parameters managed by this hook manager.

        This method collects parameters from all hooks and returns them as a flat list.

        Returns:
            List[nn.Parameter]: List of parameters from all hooks
        """
        return self._collect_from_hooks("get_parameters")

    def get_trainable_parameters(
        self,
    ) -> List[nn.Parameter]:
        """
        Get all trainable parameters from hooks.

        Returns:
            List[nn.Parameter]: List of trainable parameters
        """
        return self._collect_from_hooks("get_trainable_parameters")

    # TaskAwareInterface methods
    def register_task(
        self, task_id: int, config: Optional[Dict[str, Any]] = None
    ) -> None:
        """
        Register a new task with this hook.

        Args:
            task_id: Identifier for the new task
            config: Optional task-specific configuration
        """
        if task_id not in self._registered_tasks:
            self._registered_tasks.append(task_id)
            self._registered_tasks.sort()  # Keep tasks sorted for consistency
            self._task_configs[task_id] = config

        # Propagate to all hooks
        self._apply_to_hooks("register_task", task_id, config)

    def set_task(self, task_id: int) -> None:
        """
        Configure the model for a specific task.

        Args:
            task_id: Task identifier
        """
        self._apply_to_hooks("set_task", task_id)
        self._current_task = task_id

    def set_trainable_tasks(self, task_ids: int | List[int], **kwargs) -> None:
        """
        Prepare this hook for a specific task.

        Args:
            task_id: Identifier for the task to prepare
            **kwargs: Additional preparation parameters
        """
        self._apply_to_hooks("set_trainable_tasks", task_ids)

    @property
    def current_task(self) -> int:
        """Get the current task ID."""
        return self._current_task

    @property
    def registered_tasks(self) -> List[int]:
        """Get list of registered task IDs."""
        return (
            self._registered_tasks.copy()
        )  # Return a copy to prevent external modification
